Add Python 3.12 support with Keras 3 migration - #522
Conversation
- Update TensorFlow from 2.15 to >=2.16 - Update Keras from 2.15 to >=3.0 - Add keras-nlp >=0.24.0 for Tokenizer compatibility - Migrate to Keras 3 APIs: - Replace Model.add_loss() with FeatureLossLayer for Functional models - Replace optimizer.minimize() with tape.gradient() + apply_gradients() - Update weight file extension from .ckpt to .weights.h5 - Use keras.ops instead of tf.reduce_sum/exp - Use keras.random with seed generator for reproducibility - Update Dockerfile to Python 3.12 base image - All 686 unit tests pass
There was a problem hiding this comment.
Pull request overview
This PR migrates the codebase from TensorFlow 2.15/Keras 2.15 to TensorFlow 2.16+/Keras 3.0+, adding Python 3.12 support. The migration involves significant architectural changes to accommodate Keras 3's different API patterns, particularly around loss computation in Functional models and optimizer usage. The PR updates model weight file extensions from .ckpt to .weights.h5 and introduces proper random seed management through Keras 3's SeedGenerator API.
Key changes include:
- Replacing
Model.add_loss()withFeatureLossLayerthat usesLayer.add_loss()for Keras 3 Functional model compatibility - Replacing
optimizer.minimize()with manual gradient computation viatape.gradient()+apply_gradients() - Migrating from TensorFlow-specific ops (
tf.reduce_sum,tf.exp) to Keras ops (keras.ops) for better backend flexibility
Reviewed changes
Copilot reviewed 6 out of 8 changed files in this pull request and generated 10 comments.
Show a summary per file
| File | Description |
|---|---|
| requirements.txt | Updates TensorFlow to >=2.16, Keras to >=3.0, adds keras-nlp dependency, removes xlwt |
| Dockerfile | Updates base image from Python 3.11 to Python 3.12 |
| src/syngen/ml/vae/models/custom_layers.py | Implements FeatureLossLayer for Keras 3 Functional models, adds module-level seed generator, migrates to keras.ops and keras.random APIs |
| src/syngen/ml/vae/models/model.py | Refactors model building to use FeatureLossLayer instead of Model.add_loss(), removes KL loss computation, updates imports to keras namespace, replaces Activation with LeakyReLU layers |
| src/syngen/ml/vae/models/features.py | Migrates to keras namespace imports, replaces K.random_uniform_variable with keras.random.uniform using seed generator, updates loss functions to use keras.ops |
| src/syngen/ml/vae/wrappers/wrappers.py | Replaces optimizer.minimize() with manual gradient application, updates weight file extensions to .weights.h5, simplifies optimizer creation, adds random_seed parameter |
| src/syngen/ml/handlers/handlers.py | Adds random_seed parameter propagation to wrapper initialization |
| .gitignore | Adds test_data/ directory to ignore list |
Comments suppressed due to low confidence (1)
src/syngen/ml/vae/models/model.py:72
- The
_create_feature_loss_layermethod is no longer used after the Keras 3 migration. It was previously called to create loss layers, but now FeatureLossLayer is instantiated directly in thebuild_modelmethod (lines 129-134). Consider removing this unused method to reduce code clutter.
@staticmethod
@slugify_parameters(exclude_params=("feature",))
def _create_feature_loss_layer(feature, name):
FeatureLossLayer(feature, name=name)
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # KL loss is disabled (weight 0) in original code | ||
| kl_loss = 0.0 |
There was a problem hiding this comment.
The KL loss is hardcoded to 0.0 with a comment explaining it's disabled. However, this doesn't match the original implementation where KL loss was computed (even if multiplied by 0). If the KL loss computation was removed from the model entirely, the comment and this line are correct. However, if there's any chance of re-enabling KL loss in the future, this hardcoded value removes that flexibility.
Consider either:
- Removing the KL loss tracking entirely if it's permanently disabled
- Computing it properly but with a weight of 0, maintaining the ability to re-enable it via configuration
| # Compute gradients and apply them | ||
| gradients = tape.gradient(loss, self.model.trainable_weights) | ||
| self.optimizer.apply_gradients(zip(gradients, self.model.trainable_weights)) |
There was a problem hiding this comment.
The gradient computation doesn't check for None gradients. If the loss is not properly connected to some trainable variables, tape.gradient() may return None for those gradients, causing apply_gradients to fail.
Consider adding a check:
gradients = tape.gradient(loss, self.model.trainable_weights)
if gradients is None or None in gradients:
raise ValueError("Gradients contain None - check loss computation")
self.optimizer.apply_gradients(zip(gradients, self.model.trainable_weights))There was a problem hiding this comment.
@copilot open a new pull request to apply changes based on this feedback
| # Compute loss based on feature type | ||
| if self.loss_type == 'continuous': | ||
| loss = self.weight * ops.mean(keras.losses.mean_squared_error(feature_input, feature_decoder)) | ||
| elif self.loss_type == 'binary': | ||
| loss = self.weight * ops.mean(keras.losses.binary_crossentropy(feature_input, feature_decoder)) | ||
| else: # categorical | ||
| loss = self.weight * ops.mean(keras.losses.categorical_crossentropy(feature_input, feature_decoder)) | ||
|
|
||
| self.add_loss(loss) |
There was a problem hiding this comment.
The FeatureLossLayer computes a simplified loss and ignores the original feature.loss property, which includes custom logic like random weight sampling. This changes the training behavior from the original implementation where each feature could have its own loss function with weight_randomizer support.
The original code used feature.loss which could include random weights (e.g., random_weight * keras.losses.mean_squared_error(...) in ContinuousFeature.loss). The new implementation uses a fixed weight parameter, losing the dynamic weight randomization functionality.
Consider either:
- Using
feature.lossdirectly in the FeatureLossLayer, or - Passing the
weight_randomizerto FeatureLossLayer and implementing the random weight logic there.
| # Module-level seed generator for reproducible random operations | ||
| _seed_generator: Optional[keras.random.SeedGenerator] = None | ||
|
|
||
|
|
||
| def set_seed_generator(seed: Optional[int] = None): | ||
| """ | ||
| Set the module-level seed generator for reproducible random operations. | ||
| Call this before building the VAE model. | ||
| """ | ||
| global _seed_generator | ||
| if seed is not None: | ||
| _seed_generator = keras.random.SeedGenerator(seed) | ||
| else: | ||
| _seed_generator = None |
There was a problem hiding this comment.
The module-level _seed_generator is a global mutable state that could cause issues in multi-threaded or multi-process environments. If multiple VAE instances are created with different seeds in parallel, they will overwrite each other's seed generator.
Consider making the seed generator an instance variable of the CVAE class instead of a module-level global, or use thread-local storage if parallelization is required.
| @@ -1,14 +1,14 @@ | |||
| # syntax=docker/dockerfile:1 | |||
There was a problem hiding this comment.
[nitpick] The PR description mentions uncertainty about whether to keep this in a separate branch for Python 3.12 ("maybe we need to create new dockerfile for 3.12"). Consider adding a comment in the Dockerfile indicating this is for Python 3.12 with Keras 3, or creating a separate Dockerfile (e.g., Dockerfile.python3.12) as suggested in the PR description, maintaining backward compatibility with Python 3.11.
| # syntax=docker/dockerfile:1 | |
| # syntax=docker/dockerfile:1 | |
| # This Dockerfile is for Python 3.12 with Keras 3. |
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
- Update setup.cfg files to allow Python 3.12 - Fix Keras 3 API compatibility: FeatureLossLayer custom_loss parameter - Fix Keras 3 API compatibility: categorical_crossentropy parameter names
Сheck as it may worse to hold it in separate branch for now and build some 3.12 version separately as experimental or something